import json
import csv
import logging
import os
import re
from collections import deque

from sonetto.agents.base import BaseAgent, ContextLevel
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.message import Message
from sonetto.core.routing.broadcast import CommonBroadcastConditions
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset, field

base_url_glm = "https://open.bigmodel.cn/api/paas/v4"
api_key_glm = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model_glm = "glm-4-flash"

base_url_ds = "https://api.deepseek.com"
api_key_ds = "sk-3bb4d61950df4f4d833f800848e1f461"
model_ds = "deepseek-chat"

toolset = McpToolset()


chunker = ChatAgent(agent_id="chunker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    system_prompt='''
你是一个文档语义边界划分专家，你的任务是识别文本的语义边界。
你划分的块最后会作为问答系统的参考内容，故其语义的独立完整性非常重要。

### 语义边界划分指南：
1. 语义边界之间的块应在主题/语义上连贯且独立。
2. 语义边界划分：
    - 分节边界（章节、标题）
    - 内容主题转换
    - 形成逻辑单元的段落组
    - 如果看到问答对（先问再答，不能先答再问），请直接将一个问答对作为一个块——一个块只能包含一个问答对！
    - 有问必须有答，不能出现有问无答的分块！
    - 分隔符这种没有明确语义信息的标志不能作为分块依据！
3. 图片或表格不要单独划分，应当和最近的文本块合并。
4. 分块粒度不宜太粗（比如分成整个章节），也不宜太细（比如按句划分）。
5. 边界行的上面部分（不含边界行）会被归为上一个块，下面部分（包含边界行）会被归为下一个块。

由于文档非常庞大，用户只能分多次为你提供文档内容，每次只能提供100行。

请你具体分析语义边界的划分原因，最后以json数组的格式给出你的语义边界行号，示例回答：

<answer>
/*此处分析每个分块的主题和语义独立完整性……*/
```json
[5, 36, 54]
```
</answer>
在上面的回答中，包括第4行的以前的内容作为一块，[5, 35] 作为一块，[36, 53] 作为一块，[54, ...] 作为一块。
''')


class CoordinatorAgent(BaseAgent):

    def __init__(self, agent_id: str, working_directory: str):
        super().__init__(agent_id)
        self.directory = working_directory
        self.document_queue = deque([
            doc for doc in os.listdir(working_directory) if doc.endswith('md')
        ])
        logging.debug(f"文档队列：{self.document_queue}")
        self.current_document = None
        self.current_output = None
        self.current_line = 0
        self.total_lines = 0
        self.window_size = 200
        self.bounds, self.chunks = [], deque()

    def list_directory(self):
        return os.listdir(self.directory)

    def pretty_print(self, content: dict[str], lineno: bool = True) -> str:
        ret = ""
        for key, value in content.items():
            if lineno:
                ret += f"第{key}行: {value}\n"
            else:
                ret += f"{value}\n"
        ret += "-----------------\n"
        return ret

    def read_lines(self, start: int, end: int, lineno: bool = True) -> str:
        """
        读取文件中指定范围的行
        
        Args:
            filename: 文件路径
            start: 起始行号（从1开始，包含）
            end: 结束行号（从1开始，包含）
            
        Returns:
            包含行号和内容的字符串列表，每个字符串格式为 "行号: 内容"
            
        Raises:
            FileNotFoundError: 文件不存在时抛出
            ValueError: 起始或结束行号无效时抛出
            TypeError: 行号不是整数时抛出
        """
        # 参数验证
        if not isinstance(start, int) or not isinstance(end, int):
            raise TypeError("行号必须是整数")

        if start < 1:
            raise ValueError("起始行号必须大于等于1")

        if end < start:
            raise ValueError("结束行号不能小于起始行号")

        try:
            result = {}
            with open(self.current_document, 'r', encoding='utf-8') as file:
                line_num = 0
                for line in file:
                    line_num += 1

                    if line_num < start:
                        continue

                    if line_num <= end:
                        result[line_num] = f"{line.rstrip()}"
                    else:
                        break

                if not result and start > line_num:
                    raise ValueError(f"文件只有{line_num}行，少于起始行号{start}")
                self.current_line += self.window_size
                return self.pretty_print(result, lineno)

        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def split_md2jsonl(self) -> None:
        self.bounds.insert(0, 1)
        self.bounds.append(self.total_lines + 1)
        print(self.bounds)
        with open(self.current_document, 'r', encoding='utf-8') as file:
            with open(self.current_output, 'w', encoding='utf-8') as output:
                writer = csv.writer(output,
                                    quoting=csv.QUOTE_MINIMAL,
                                    escapechar='\\')
                for i in range(1, len(self.bounds)):
                    content = self.read_lines(self.bounds[i - 1],
                                              self.bounds[i] - 1, False)
                    # output.write(
                    #     json.dumps({'content': content}, ensure_ascii=False) + '\n')
                    # writer.writerow([json.dumps(content, ensure_ascii=False)])
                    writer.writerow([content])
                    self.chunks.append(content)
        self.bounds.clear()

    def count_lines(self) -> int:
        try:
            with open(self.current_document, 'r', encoding='utf-8') as file:
                return len(file.readlines())
        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def turn2next_document(self) -> None:
        if self.document_queue:
            self.current_document = os.path.join(self.directory,
                                                 self.document_queue.popleft())
            self.current_output = f"{self.current_document}.csv"
            self.detail_output = f"{self.current_document}_propositions.csv"
            self.current_line = 1
            self.total_lines = self.count_lines()
            return True
        else:
            return False

    def parse_bounds(self, content: str) -> list[int]:
        bounds = re.findall(r"```json\n(.+?)\n```", content, re.DOTALL)
        print(bounds)
        bounds = json.loads(bounds[0])
        self.bounds.extend([
            ele for ele in bounds
            if ele % self.window_size != 0 and ele % self.window_size != 1
        ])

    def next_chunk(self) -> str:
        if self.chunks:
            return self.chunks.popleft()
        else:
            return ""

    def forward(self, message: Message) -> Message:
        sender = message.sender
        receiver = message.receiver
        content = message.content

        if sender == '__init__':
            if not self.turn2next_document():
                return Message(content="没有文档需要处理",
                               sender=self.agent_id,
                               receiver='chunker')
            else:
                content = f"""
接下来开始对文档 '{self.current_document}' 进行分块。首先需要分块的文本内容如下所示：
```markdown
{self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
```
请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
"""
                return Message(content=content,
                               sender=self.agent_id,
                               receiver='chunker')

        # 检查是否有文档需要处理
        if self.current_line < self.total_lines:
            # 提取json数组中的元素，拼接到self.bounds中
            self.parse_bounds(content)
            content = f"""
继续对文档 '{self.current_document}' 进行分块。接下来需要分块的文本内容如下所示：
```markdown
{self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
```
请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
"""
            return Message(content=content,
                           sender=self.agent_id,
                           receiver=sender)
        else:
            self.parse_bounds(content)
            chunker.clear_context(ContextLevel.TASK)
            self.split_md2jsonl()
            if not self.turn2next_document():
                    return Message(content="所有文档均已整理完毕，圆满结束！！！",
                                   sender=self.agent_id,
                                   receiver=sender)
            content = f"""
    接下来开始对文档 '{self.current_document}' 进行分块。首先需要分块的文本内容如下所示：
    ```markdown
    {self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
    ```
    请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
    """
            return Message(content=content,
                           sender=self.agent_id,
                           receiver='chunker')


coordinator = CoordinatorAgent(agent_id="coordinator",
                               working_directory=".\HSBC")


scheduler = Scheduler()
scheduler.register_agent(coordinator)
scheduler.register_agent(chunker)
scheduler.set_route(chunker.agent_id, coordinator.agent_id)

if __name__ == "__main__":
    logging.debug("文档分块任务开始")
    scheduler.run_stream("请完成文档分块任务",
                         cancel_token="圆满结束",
                         cancel_from=[coordinator.agent_id],
                         max_step=1500)
    logging.debug("文档分块任务结束")
